package com.small.xx.ai.spring.observability.config;

import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationHandler;
import io.micrometer.observation.ObservationRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.client.observation.ChatClientObservationContext;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.observation.ChatModelObservationContext;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.embedding.observation.EmbeddingModelObservationContext;
import org.springframework.ai.observation.AiOperationMetadata;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.tool.observation.ToolCallingObservationContext;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.util.List;

@Configuration
@ConditionalOnProperty(prefix = "spring.ai.observation", name = "enabled", havingValue = "true", matchIfMissing = true)
public class ObservationConfiguration {

    private static final Logger logger = LoggerFactory.getLogger(ObservationConfiguration.class);

    /**
     * 创建ObservationRegistry Bean（观测注册表），用于管理所有观测处理器
     * @ConditionalOnMissingBean 表示当容器中没有名为"observationRegistry"的Bean时才创建
     */
    @Bean
    @ConditionalOnMissingBean(name = "observationRegistry")
    public ObservationRegistry observationRegistry(
        // 注入所有ObservationHandler类型的Bean（通过ObjectProvider实现延迟注入）
        ObjectProvider<ObservationHandler<?>> observationHandlerObjectProvider) {
        // 创建一个新的观测注册表实例
        ObservationRegistry observationRegistry = ObservationRegistry.create();
        // 获取注册表的配置对象，用于添加观测处理器
        ObservationRegistry.ObservationConfig observationConfig = observationRegistry.observationConfig();

        // 遍历所有注入的观测处理器（按顺序）
        observationHandlerObjectProvider.orderedStream().forEach(handler -> {
            // 获取处理器类实现的所有泛型接口
            Type[] genericInterfaces = handler.getClass().getGenericInterfaces();
            // 遍历每个接口，解析其泛型参数（判断支持的上下文类型）
            for (Type type : genericInterfaces) {
                // 检查是否为参数化类型（带泛型的接口）
                if (type instanceof ParameterizedType parameterizedType //
                    // 检查原始类型是否为ObservationHandler的实现类
                    && parameterizedType.getRawType() instanceof Class<?> clazz //
                    && ObservationHandler.class.isAssignableFrom(clazz)//
                ) {
                    // 获取泛型接口的实际类型参数（即该处理器支持的上下文类型）
                    Type actualTypeArgument = parameterizedType.getActualTypeArguments()[0];
                    // 日志记录：加载的处理器支持的上下文类型
                    logger.info("加载观察处理程序，支持上下文类型: {}", actualTypeArgument);
                }
            }

            // 将当前处理器添加到观测注册表的配置中（使其生效）
            observationConfig.observationHandler(handler);
        });
        // 返回配置完成的观测注册表
        return observationRegistry;
    }

    /**
     * 创建监听ChatClient调用的观测处理器Bean 用于监控ChatClient（聊天客户端）的调用过程
     */
    @Bean
    ObservationHandler<ChatClientObservationContext> chatClientObservationContextObservationHandler() {
        // 日志记录：ChatClient观测器启动
        logger.info("ChatClientObservation 启动");
        // 返回一个匿名实现类，处理ChatClient相关的观测事件
        return new ObservationHandler<>() {

            // 判断当前处理器是否支持指定的上下文类型
            @Override
            public boolean supportsContext(Observation.Context context) {
                // 仅支持ChatClientObservationContext类型的上下文
                return context instanceof ChatClientObservationContext;
            }

            // 当观测对象启动时触发（ChatClient开始调用前）
            @Override
            public void onStart(ChatClientObservationContext context) {
                // 从上下文中获取ChatClient的请求对象
                ChatClientRequest request = context.getRequest();
                // 获取请求中包含的增强器（Advisors）列表
                List<? extends Advisor> advisors = context.getAdvisors();
                // 判断是否为流式调用
                boolean stream = context.isStream();
                // 日志记录：ChatClient启动时的详细信息（请求内容、增强器、是否流式）
                logger.info("💬ChatClientObservation 启动: ChatClientRequest : {}, Advisors : {}, stream : {}", request,
                    advisors, stream);
            }

            // 当观测对象停止时触发（ChatClient调用结束后）
            @Override
            public void onStop(ChatClientObservationContext context) {
                // 调用父类默认实现（可根据需要重写自定义逻辑）
                ObservationHandler.super.onStop(context);
            }
        };
    }

    /**
     * 创建监听ChatModel调用的观测处理器Bean 用于监控ChatModel（聊天模型）的调用过程
     */
    @Bean
    ObservationHandler<ChatModelObservationContext> chatModelObservationContextObservationHandler() {
        // 日志记录：ChatModel观测器启动
        logger.info("ChatModelObservation start");
        // 返回一个匿名实现类，处理ChatModel相关的观测事件
        return new ObservationHandler<>() {

            // 判断当前处理器是否支持指定的上下文类型
            @Override
            public boolean supportsContext(Observation.Context context) {
                // 仅支持ChatModelObservationContext类型的上下文
                return context instanceof ChatModelObservationContext;
            }

            // 当观测对象启动时触发（ChatModel开始调用前）
            @Override
            public void onStart(ChatModelObservationContext context) {
                // 获取模型操作的元数据（如操作类型、提供商等）
                AiOperationMetadata operationMetadata = context.getOperationMetadata();
                // 获取模型的输入提示（Prompt）
                Prompt request = context.getRequest();
                // 日志记录：ChatModel启动时的元数据和输入提示
                logger.info("🤖ChatModelObservation 启动: AiOperationMetadata : {}", operationMetadata);
                logger.info("🤖ChatModelObservation 启动: Prompt : {}", request);
            }

            // 当观测对象停止时触发（ChatModel调用结束后）
            @Override
            public void onStop(ChatModelObservationContext context) {
                // 获取模型的输出响应
                ChatResponse response = context.getResponse();
                // 日志记录：ChatModel返回的响应内容
                logger.info("🤖ChatModelObservation 启动: ChatResponse : {}", response);
            }
        };
    }

    /**
     * 创建监听工具调用的观测处理器Bean 用于监控工具调用（如外部API、数据库操作等）的过程
     */
    @Bean
    public ObservationHandler<ToolCallingObservationContext> toolCallingObservationContextObservationHandler() {
        // 日志记录：工具调用观测器启动
        logger.info("ToolCallingObservation start");
        // 返回一个匿名实现类，处理工具调用相关的观测事件
        return new ObservationHandler<>() {
            // 判断当前处理器是否支持指定的上下文类型
            @Override
            public boolean supportsContext(Observation.Context context) {
                // 仅支持ToolCallingObservationContext类型的上下文
                return context instanceof ToolCallingObservationContext;
            }

            // 当观测对象启动时触发（工具调用开始前）
            @Override
            public void onStart(ToolCallingObservationContext context) {
                // 获取工具的定义信息（如工具名称、描述等）
                ToolDefinition toolDefinition = context.getToolDefinition();
                // 日志记录：工具调用开始，包含工具名称和输入参数
                logger.info("🔨ToolCalling start: {} - {}", toolDefinition.name(), context.getToolCallArguments());
            }

            // 当观测对象停止时触发（工具调用结束后）
            @Override
            public void onStop(ToolCallingObservationContext context) {
                // 获取工具的定义信息
                ToolDefinition toolDefinition = context.getToolDefinition();
                // 日志记录：工具调用完成，包含工具名称和返回结果
                logger.info("✅ToolCalling done: {} - {}", toolDefinition.name(), context.getToolCallResult());
            }
        };
    }

    /**
     * 创建监听EmbeddingModel调用的观测处理器Bean 用于监控嵌入模型（将文本转换为向量）的调用过程
     */
    @Bean
    public ObservationHandler<EmbeddingModelObservationContext> embeddingModelObservationContextObservationHandler() {
        // 日志记录：嵌入模型观测器启动
        logger.info("EmbeddingModelObservation start");
        // 返回一个匿名实现类，处理嵌入模型相关的观测事件
        return new ObservationHandler<>() {
            // 判断当前处理器是否支持指定的上下文类型
            @Override
            public boolean supportsContext(Observation.Context context) {
                // 仅支持EmbeddingModelObservationContext类型的上下文
                return context instanceof EmbeddingModelObservationContext;
            }

            // 当观测对象启动时触发（嵌入模型开始调用前）
            @Override
            public void onStart(EmbeddingModelObservationContext context) {
                // 日志记录：嵌入模型启动，包含操作类型和提供商信息
                logger.info("📚EmbeddingModelObservation start: {} - {}", context.getOperationMetadata().operationType(),
                    context.getOperationMetadata().provider());
            }
        };
    }
}
